{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Execute this cell to install dependencies\n",
    "%pip install sf-hamilton[visualization]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Feature catalog example for spark  [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/dagworks-inc/hamilton/blob/main/examples/spark/pyspark_feature_catalog/example_usage.ipynb) [![GitHub badge](https://img.shields.io/badge/github-view_source-2b3137?logo=github)](https://github.com/apache/hamilton/blob/main/examples/spark/pyspark_feature_catalog/example_usage.ipynb)\n",
    "\n",
    "\n",
    "Define multiple features with clear lineage using Hamilton.\n",
    "Use a custom output builder to join the features of interest together before returning them to the user.\n",
    "\n",
    "As user you only need to specify which features (groups of features) you want to compute.\n",
    "\n",
    "## Including column level lineage\n",
    "By using the `@schema.output` decorator you can extend the lineage up to column level, e.g.:\n",
    "\n",
    "```python\n",
    "@schema.output((\"zone\", \"str\"), (\"level\", \"int\"), (\"avatarId\", \"int\"))\n",
    "```\n",
    "\n",
    "## Download data\n",
    "To get started first download the data: \n",
    "```\n",
    "curl https://storage.googleapis.com/shareddatasets/wow.parquet -o data/wow.parquet\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-05-22T22:27:16.952106Z",
     "start_time": "2023-05-22T22:27:15.116241Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/elijahbenizzy/.pyenv/versions/3.11.4/envs/hamilton-3-11-fresh/lib/python3.11/site-packages/pyspark/pandas/__init__.py:50: UserWarning: 'PYARROW_IGNORE_TIMEZONE' environment variable was not set. It is required to set this environment variable to '1' in both driver and executor sides if you use pyarrow>=2.0.0. pandas-on-Spark will set it for you but it does not work if there is a Spark context already launched.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import findspark\n",
    "findspark.init()\n",
    "\n",
    "import importlib\n",
    "import logging\n",
    "import sys\n",
    "import typing\n",
    "\n",
    "import pyspark.sql as ps\n",
    "\n",
    "from hamilton import base\n",
    "from hamilton import driver\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout)\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a customer builder and adapter to create the final output (spark DataFrame)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyCustomBuilder(base.ResultMixin):\n",
    "     @staticmethod\n",
    "     def build_result(**outputs: typing.Dict[str, typing.Any]) -> ps.DataFrame:\n",
    "        # TODO: add error handling when incompatible outputs are created\n",
    "        level_info = outputs['level_info']\n",
    "        zone_counts = outputs['zone_counts']\n",
    "        return zone_counts.join(level_info, on=aggregation_level)\n",
    "\n",
    "adapter = base.SimplePythonGraphAdapter(MyCustomBuilder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create the driver and show the lineage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:hamilton.telemetry:Note: Hamilton collects completely anonymous data about usage. This will help us improve Hamilton over time. See https://github.com/apache/hamilton#usage-analytics--data-privacy for details.\n"
     ]
    },
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 8.1.0 (20230707.0739)\n",
       " -->\n",
       "<!-- Pages: 1 -->\n",
       "<svg width=\"1068pt\" height=\"511pt\"\n",
       " viewBox=\"0.00 0.00 1068.35 511.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 507)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-507 1064.35,-507 1064.35,4 -4,4\"/>\n",
       "<g id=\"clust4\" class=\"cluster\">\n",
       "<title>cluster_world_of_warcraft</title>\n",
       "<path fill=\"#ffffff\" stroke=\"#649fb3\" stroke-dasharray=\"5,2\" d=\"M155.23,-8C155.23,-8 290.58,-8 290.58,-8 296.58,-8 302.58,-14 302.58,-20 302.58,-20 302.58,-242 302.58,-242 302.58,-248 296.58,-254 290.58,-254 290.58,-254 155.23,-254 155.23,-254 149.23,-254 143.23,-248 143.23,-242 143.23,-242 143.23,-20 143.23,-20 143.23,-14 149.23,-8 155.23,-8\"/>\n",
       "</g>\n",
       "<g id=\"clust5\" class=\"cluster\">\n",
       "<title>cluster__legend</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"15.25,-254 15.25,-495 100.1,-495 100.1,-254 15.25,-254\"/>\n",
       "<text text-anchor=\"middle\" x=\"57.67\" y=\"-477.7\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">Legend</text>\n",
       "</g>\n",
       "<g id=\"clust1\" class=\"cluster\">\n",
       "<title>cluster_zone_counts</title>\n",
       "<path fill=\"#ffffff\" stroke=\"#649fb3\" stroke-dasharray=\"5,2\" d=\"M908.75,-153C908.75,-153 1040.35,-153 1040.35,-153 1046.35,-153 1052.35,-159 1052.35,-165 1052.35,-165 1052.35,-387 1052.35,-387 1052.35,-393 1046.35,-399 1040.35,-399 1040.35,-399 908.75,-399 908.75,-399 902.75,-399 896.75,-393 896.75,-387 896.75,-387 896.75,-165 896.75,-165 896.75,-159 902.75,-153 908.75,-153\"/>\n",
       "</g>\n",
       "<g id=\"clust2\" class=\"cluster\">\n",
       "<title>cluster_level_info</title>\n",
       "<path fill=\"#ffffff\" stroke=\"#649fb3\" stroke-dasharray=\"5,2\" d=\"M368.08,-232C368.08,-232 481.68,-232 481.68,-232 487.68,-232 493.68,-238 493.68,-244 493.68,-244 493.68,-358 493.68,-358 493.68,-364 487.68,-370 481.68,-370 481.68,-370 368.08,-370 368.08,-370 362.08,-370 356.08,-364 356.08,-358 356.08,-358 356.08,-244 356.08,-244 356.08,-238 362.08,-232 368.08,-232\"/>\n",
       "</g>\n",
       "<g id=\"clust3\" class=\"cluster\">\n",
       "<title>cluster_with_flags</title>\n",
       "<path fill=\"#ffffff\" stroke=\"#649fb3\" stroke-dasharray=\"5,2\" d=\"M750.15,-12C750.15,-12 875.75,-12 875.75,-12 881.75,-12 887.75,-18 887.75,-24 887.75,-24 887.75,-354 887.75,-354 887.75,-360 881.75,-366 875.75,-366 875.75,-366 750.15,-366 750.15,-366 744.15,-366 738.15,-360 738.15,-354 738.15,-354 738.15,-24 738.15,-24 738.15,-18 744.15,-12 750.15,-12\"/>\n",
       "</g>\n",
       "<!-- zone_counts -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>zone_counts</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M1014.6,-388.8C1014.6,-388.8 934.5,-388.8 934.5,-388.8 928.5,-388.8 922.5,-382.8 922.5,-376.8 922.5,-376.8 922.5,-337.2 922.5,-337.2 922.5,-331.2 928.5,-325.2 934.5,-325.2 934.5,-325.2 1014.6,-325.2 1014.6,-325.2 1020.6,-325.2 1026.6,-331.2 1026.6,-337.2 1026.6,-337.2 1026.6,-376.8 1026.6,-376.8 1026.6,-382.8 1020.6,-388.8 1014.6,-388.8\"/>\n",
       "<text text-anchor=\"start\" x=\"933.3\" y=\"-365.7\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">zone_counts</text>\n",
       "<text text-anchor=\"start\" x=\"940.43\" y=\"-337.7\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">DataFrame</text>\n",
       "</g>\n",
       "<!-- total_count=int &#160;&#160;&#160; -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>total_count=int &#160;&#160;&#160;</title>\n",
       "<path fill=\"#c8dae0\" stroke=\"black\" d=\"M1030.35,-307C1030.35,-307 918.75,-307 918.75,-307 912.75,-307 906.75,-301 906.75,-295 906.75,-295 906.75,-283 906.75,-283 906.75,-277 912.75,-271 918.75,-271 918.75,-271 1030.35,-271 1030.35,-271 1036.35,-271 1042.35,-277 1042.35,-283 1042.35,-283 1042.35,-295 1042.35,-295 1042.35,-301 1036.35,-307 1030.35,-307\"/>\n",
       "<text text-anchor=\"middle\" x=\"974.55\" y=\"-285.88\" font-family=\"Courier,monospace\" font-size=\"10.00\">total_count=int &#160;&#160;&#160;</text>\n",
       "</g>\n",
       "<!-- darkshore_count=int -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>darkshore_count=int</title>\n",
       "<path fill=\"#c8dae0\" stroke=\"black\" d=\"M1030.35,-253C1030.35,-253 918.75,-253 918.75,-253 912.75,-253 906.75,-247 906.75,-241 906.75,-241 906.75,-229 906.75,-229 906.75,-223 912.75,-217 918.75,-217 918.75,-217 1030.35,-217 1030.35,-217 1036.35,-217 1042.35,-223 1042.35,-229 1042.35,-229 1042.35,-241 1042.35,-241 1042.35,-247 1036.35,-253 1030.35,-253\"/>\n",
       "<text text-anchor=\"middle\" x=\"974.55\" y=\"-231.88\" font-family=\"Courier,monospace\" font-size=\"10.00\">darkshore_count=int</text>\n",
       "</g>\n",
       "<!-- durotar_count=int &#160; -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>durotar_count=int &#160;</title>\n",
       "<path fill=\"#c8dae0\" stroke=\"black\" d=\"M1030.35,-199C1030.35,-199 918.75,-199 918.75,-199 912.75,-199 906.75,-193 906.75,-187 906.75,-187 906.75,-175 906.75,-175 906.75,-169 912.75,-163 918.75,-163 918.75,-163 1030.35,-163 1030.35,-163 1036.35,-163 1042.35,-169 1042.35,-175 1042.35,-175 1042.35,-187 1042.35,-187 1042.35,-193 1036.35,-199 1030.35,-199\"/>\n",
       "<text text-anchor=\"middle\" x=\"974.55\" y=\"-177.88\" font-family=\"Courier,monospace\" font-size=\"10.00\">durotar_count=int &#160;</text>\n",
       "</g>\n",
       "<!-- spark_session -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>spark_session</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M103.35,-243.8C103.35,-243.8 12,-243.8 12,-243.8 6,-243.8 0,-237.8 0,-231.8 0,-231.8 0,-192.2 0,-192.2 0,-186.2 6,-180.2 12,-180.2 12,-180.2 103.35,-180.2 103.35,-180.2 109.35,-180.2 115.35,-186.2 115.35,-192.2 115.35,-192.2 115.35,-231.8 115.35,-231.8 115.35,-237.8 109.35,-243.8 103.35,-243.8\"/>\n",
       "<text text-anchor=\"start\" x=\"10.8\" y=\"-220.7\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">spark_session</text>\n",
       "<text text-anchor=\"start\" x=\"16.05\" y=\"-192.7\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">SparkSession</text>\n",
       "</g>\n",
       "<!-- world_of_warcraft -->\n",
       "<g id=\"node16\" class=\"node\">\n",
       "<title>world_of_warcraft</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M280.58,-243.8C280.58,-243.8 165.23,-243.8 165.23,-243.8 159.23,-243.8 153.23,-237.8 153.23,-231.8 153.23,-231.8 153.23,-192.2 153.23,-192.2 153.23,-186.2 159.23,-180.2 165.23,-180.2 165.23,-180.2 280.58,-180.2 280.58,-180.2 286.58,-180.2 292.58,-186.2 292.58,-192.2 292.58,-192.2 292.58,-231.8 292.58,-231.8 292.58,-237.8 286.58,-243.8 280.58,-243.8\"/>\n",
       "<text text-anchor=\"start\" x=\"164.03\" y=\"-220.7\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">world_of_warcraft</text>\n",
       "<text text-anchor=\"start\" x=\"188.78\" y=\"-192.7\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">DataFrame</text>\n",
       "</g>\n",
       "<!-- spark_session&#45;&gt;world_of_warcraft -->\n",
       "<g id=\"edge9\" class=\"edge\">\n",
       "<title>spark_session&#45;&gt;world_of_warcraft</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M115.76,-212C124.26,-212 133.15,-212 142.01,-212\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"141.77,-215.5 151.77,-212 141.77,-208.5 141.77,-215.5\"/>\n",
       "</g>\n",
       "<!-- with_flags.durotar_flag -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>with_flags.durotar_flag</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M707.15,-243.8C707.15,-243.8 560.3,-243.8 560.3,-243.8 554.3,-243.8 548.3,-237.8 548.3,-231.8 548.3,-231.8 548.3,-192.2 548.3,-192.2 548.3,-186.2 554.3,-180.2 560.3,-180.2 560.3,-180.2 707.15,-180.2 707.15,-180.2 713.15,-180.2 719.15,-186.2 719.15,-192.2 719.15,-192.2 719.15,-231.8 719.15,-231.8 719.15,-237.8 713.15,-243.8 707.15,-243.8\"/>\n",
       "<text text-anchor=\"start\" x=\"559.1\" y=\"-220.7\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">with_flags.durotar_flag</text>\n",
       "<text text-anchor=\"start\" x=\"599.6\" y=\"-192.7\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">DataFrame</text>\n",
       "</g>\n",
       "<!-- with_flags -->\n",
       "<g id=\"node10\" class=\"node\">\n",
       "<title>with_flags</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M845.88,-355.8C845.88,-355.8 780.03,-355.8 780.03,-355.8 774.03,-355.8 768.03,-349.8 768.03,-343.8 768.03,-343.8 768.03,-304.2 768.03,-304.2 768.03,-298.2 774.03,-292.2 780.03,-292.2 780.03,-292.2 845.88,-292.2 845.88,-292.2 851.88,-292.2 857.88,-298.2 857.88,-304.2 857.88,-304.2 857.88,-343.8 857.88,-343.8 857.88,-349.8 851.88,-355.8 845.88,-355.8\"/>\n",
       "<text text-anchor=\"start\" x=\"779.58\" y=\"-332.7\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">with_flags</text>\n",
       "<text text-anchor=\"start\" x=\"778.83\" y=\"-304.7\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">DataFrame</text>\n",
       "</g>\n",
       "<!-- with_flags.durotar_flag&#45;&gt;with_flags -->\n",
       "<g id=\"edge8\" class=\"edge\">\n",
       "<title>with_flags.durotar_flag&#45;&gt;with_flags</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M679.68,-244.22C697.61,-256.72 718.61,-270.9 738.15,-283 744.66,-287.03 751.59,-291.12 758.49,-295.09\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"756.42,-298.5 766.85,-300.39 759.88,-292.42 756.42,-298.5\"/>\n",
       "</g>\n",
       "<!-- with_flags.darkshore_flag -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>with_flags.darkshore_flag</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M507.3,-183.8C507.3,-183.8 342.45,-183.8 342.45,-183.8 336.45,-183.8 330.45,-177.8 330.45,-171.8 330.45,-171.8 330.45,-132.2 330.45,-132.2 330.45,-126.2 336.45,-120.2 342.45,-120.2 342.45,-120.2 507.3,-120.2 507.3,-120.2 513.3,-120.2 519.3,-126.2 519.3,-132.2 519.3,-132.2 519.3,-171.8 519.3,-171.8 519.3,-177.8 513.3,-183.8 507.3,-183.8\"/>\n",
       "<text text-anchor=\"start\" x=\"341.25\" y=\"-160.7\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">with_flags.darkshore_flag</text>\n",
       "<text text-anchor=\"start\" x=\"390.75\" y=\"-132.7\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">DataFrame</text>\n",
       "</g>\n",
       "<!-- with_flags.darkshore_flag&#45;&gt;with_flags.durotar_flag -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>with_flags.darkshore_flag&#45;&gt;with_flags.durotar_flag</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M519.62,-179.19C525.5,-180.9 531.42,-182.61 537.3,-184.32\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"536.27,-187.95 546.85,-187.38 538.22,-181.23 536.27,-187.95\"/>\n",
       "</g>\n",
       "<!-- level_info -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>level_info</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M457.8,-359.8C457.8,-359.8 391.95,-359.8 391.95,-359.8 385.95,-359.8 379.95,-353.8 379.95,-347.8 379.95,-347.8 379.95,-308.2 379.95,-308.2 379.95,-302.2 385.95,-296.2 391.95,-296.2 391.95,-296.2 457.8,-296.2 457.8,-296.2 463.8,-296.2 469.8,-302.2 469.8,-308.2 469.8,-308.2 469.8,-347.8 469.8,-347.8 469.8,-353.8 463.8,-359.8 457.8,-359.8\"/>\n",
       "<text text-anchor=\"start\" x=\"393.75\" y=\"-336.7\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">level_info</text>\n",
       "<text text-anchor=\"start\" x=\"390.75\" y=\"-308.7\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">DataFrame</text>\n",
       "</g>\n",
       "<!-- mean_level=float -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>mean_level=float</title>\n",
       "<path fill=\"#c8dae0\" stroke=\"black\" d=\"M471.68,-278C471.68,-278 378.08,-278 378.08,-278 372.08,-278 366.08,-272 366.08,-266 366.08,-266 366.08,-254 366.08,-254 366.08,-248 372.08,-242 378.08,-242 378.08,-242 471.68,-242 471.68,-242 477.68,-242 483.68,-248 483.68,-254 483.68,-254 483.68,-266 483.68,-266 483.68,-272 477.68,-278 471.68,-278\"/>\n",
       "<text text-anchor=\"middle\" x=\"424.88\" y=\"-256.88\" font-family=\"Courier,monospace\" font-size=\"10.00\">mean_level=float</text>\n",
       "</g>\n",
       "<!-- with_flags&#45;&gt;zone_counts -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>with_flags&#45;&gt;zone_counts</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M857.98,-333.1C874.49,-336.52 893.57,-340.46 911.41,-344.15\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"910.53,-347.75 921.03,-346.35 911.95,-340.89 910.53,-347.75\"/>\n",
       "</g>\n",
       "<!-- zone=str &#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160; -->\n",
       "<g id=\"node11\" class=\"node\">\n",
       "<title>zone=str &#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;</title>\n",
       "<path fill=\"#c8dae0\" stroke=\"black\" d=\"M865.75,-274C865.75,-274 760.15,-274 760.15,-274 754.15,-274 748.15,-268 748.15,-262 748.15,-262 748.15,-250 748.15,-250 748.15,-244 754.15,-238 760.15,-238 760.15,-238 865.75,-238 865.75,-238 871.75,-238 877.75,-244 877.75,-250 877.75,-250 877.75,-262 877.75,-262 877.75,-268 871.75,-274 865.75,-274\"/>\n",
       "<text text-anchor=\"middle\" x=\"812.95\" y=\"-252.88\" font-family=\"Courier,monospace\" font-size=\"10.00\">zone=str &#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;</text>\n",
       "</g>\n",
       "<!-- level=int &#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160; -->\n",
       "<g id=\"node12\" class=\"node\">\n",
       "<title>level=int &#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;</title>\n",
       "<path fill=\"#c8dae0\" stroke=\"black\" d=\"M865.75,-220C865.75,-220 760.15,-220 760.15,-220 754.15,-220 748.15,-214 748.15,-208 748.15,-208 748.15,-196 748.15,-196 748.15,-190 754.15,-184 760.15,-184 760.15,-184 865.75,-184 865.75,-184 871.75,-184 877.75,-190 877.75,-196 877.75,-196 877.75,-208 877.75,-208 877.75,-214 871.75,-220 865.75,-220\"/>\n",
       "<text text-anchor=\"middle\" x=\"812.95\" y=\"-198.88\" font-family=\"Courier,monospace\" font-size=\"10.00\">level=int &#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;</text>\n",
       "</g>\n",
       "<!-- avatarId=int &#160;&#160;&#160;&#160;&#160; -->\n",
       "<g id=\"node13\" class=\"node\">\n",
       "<title>avatarId=int &#160;&#160;&#160;&#160;&#160;</title>\n",
       "<path fill=\"#c8dae0\" stroke=\"black\" d=\"M865.75,-166C865.75,-166 760.15,-166 760.15,-166 754.15,-166 748.15,-160 748.15,-154 748.15,-154 748.15,-142 748.15,-142 748.15,-136 754.15,-130 760.15,-130 760.15,-130 865.75,-130 865.75,-130 871.75,-130 877.75,-136 877.75,-142 877.75,-142 877.75,-154 877.75,-154 877.75,-160 871.75,-166 865.75,-166\"/>\n",
       "<text text-anchor=\"middle\" x=\"812.95\" y=\"-144.88\" font-family=\"Courier,monospace\" font-size=\"10.00\">avatarId=int &#160;&#160;&#160;&#160;&#160;</text>\n",
       "</g>\n",
       "<!-- darkshore_flag=int -->\n",
       "<g id=\"node14\" class=\"node\">\n",
       "<title>darkshore_flag=int</title>\n",
       "<path fill=\"#c8dae0\" stroke=\"black\" d=\"M865.75,-112C865.75,-112 760.15,-112 760.15,-112 754.15,-112 748.15,-106 748.15,-100 748.15,-100 748.15,-88 748.15,-88 748.15,-82 754.15,-76 760.15,-76 760.15,-76 865.75,-76 865.75,-76 871.75,-76 877.75,-82 877.75,-88 877.75,-88 877.75,-100 877.75,-100 877.75,-106 871.75,-112 865.75,-112\"/>\n",
       "<text text-anchor=\"middle\" x=\"812.95\" y=\"-90.88\" font-family=\"Courier,monospace\" font-size=\"10.00\">darkshore_flag=int</text>\n",
       "</g>\n",
       "<!-- durotar_flag=int &#160; -->\n",
       "<g id=\"node15\" class=\"node\">\n",
       "<title>durotar_flag=int &#160;</title>\n",
       "<path fill=\"#c8dae0\" stroke=\"black\" d=\"M865.75,-58C865.75,-58 760.15,-58 760.15,-58 754.15,-58 748.15,-52 748.15,-46 748.15,-46 748.15,-34 748.15,-34 748.15,-28 754.15,-22 760.15,-22 760.15,-22 865.75,-22 865.75,-22 871.75,-22 877.75,-28 877.75,-34 877.75,-34 877.75,-46 877.75,-46 877.75,-52 871.75,-58 865.75,-58\"/>\n",
       "<text text-anchor=\"middle\" x=\"812.95\" y=\"-36.88\" font-family=\"Courier,monospace\" font-size=\"10.00\">durotar_flag=int &#160;</text>\n",
       "</g>\n",
       "<!-- world_of_warcraft&#45;&gt;with_flags.durotar_flag -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>world_of_warcraft&#45;&gt;with_flags.durotar_flag</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M292.79,-212C359.81,-212 462.21,-212 537.1,-212\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"536.95,-215.5 546.95,-212 536.95,-208.5 536.95,-215.5\"/>\n",
       "</g>\n",
       "<!-- world_of_warcraft&#45;&gt;with_flags.darkshore_flag -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>world_of_warcraft&#45;&gt;with_flags.darkshore_flag</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M293,-191.27C301.69,-188.66 310.73,-185.95 319.82,-183.22\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"320.51,-186.37 329.08,-180.14 318.5,-179.66 320.51,-186.37\"/>\n",
       "</g>\n",
       "<!-- world_of_warcraft&#45;&gt;level_info -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>world_of_warcraft&#45;&gt;level_info</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M265.4,-244.12C284.61,-258.22 308.09,-274.39 330.45,-287 342.87,-294 356.71,-300.65 369.86,-306.48\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"368.17,-310 378.74,-310.77 370.96,-303.58 368.17,-310\"/>\n",
       "</g>\n",
       "<!-- zone=str &#160;&#160;&#160; -->\n",
       "<g id=\"node17\" class=\"node\">\n",
       "<title>zone=str &#160;&#160;&#160;</title>\n",
       "<path fill=\"#c8dae0\" stroke=\"black\" d=\"M257.7,-162C257.7,-162 188.1,-162 188.1,-162 182.1,-162 176.1,-156 176.1,-150 176.1,-150 176.1,-138 176.1,-138 176.1,-132 182.1,-126 188.1,-126 188.1,-126 257.7,-126 257.7,-126 263.7,-126 269.7,-132 269.7,-138 269.7,-138 269.7,-150 269.7,-150 269.7,-156 263.7,-162 257.7,-162\"/>\n",
       "<text text-anchor=\"middle\" x=\"222.9\" y=\"-140.88\" font-family=\"Courier,monospace\" font-size=\"10.00\">zone=str &#160;&#160;&#160;</text>\n",
       "</g>\n",
       "<!-- level=int &#160;&#160; -->\n",
       "<g id=\"node18\" class=\"node\">\n",
       "<title>level=int &#160;&#160;</title>\n",
       "<path fill=\"#c8dae0\" stroke=\"black\" d=\"M257.7,-108C257.7,-108 188.1,-108 188.1,-108 182.1,-108 176.1,-102 176.1,-96 176.1,-96 176.1,-84 176.1,-84 176.1,-78 182.1,-72 188.1,-72 188.1,-72 257.7,-72 257.7,-72 263.7,-72 269.7,-78 269.7,-84 269.7,-84 269.7,-96 269.7,-96 269.7,-102 263.7,-108 257.7,-108\"/>\n",
       "<text text-anchor=\"middle\" x=\"222.9\" y=\"-86.88\" font-family=\"Courier,monospace\" font-size=\"10.00\">level=int &#160;&#160;</text>\n",
       "</g>\n",
       "<!-- avatarId=int -->\n",
       "<g id=\"node19\" class=\"node\">\n",
       "<title>avatarId=int</title>\n",
       "<path fill=\"#c8dae0\" stroke=\"black\" d=\"M257.7,-54C257.7,-54 188.1,-54 188.1,-54 182.1,-54 176.1,-48 176.1,-42 176.1,-42 176.1,-30 176.1,-30 176.1,-24 182.1,-18 188.1,-18 188.1,-18 257.7,-18 257.7,-18 263.7,-18 269.7,-24 269.7,-30 269.7,-30 269.7,-42 269.7,-42 269.7,-48 263.7,-54 257.7,-54\"/>\n",
       "<text text-anchor=\"middle\" x=\"222.9\" y=\"-32.88\" font-family=\"Courier,monospace\" font-size=\"10.00\">avatarId=int</text>\n",
       "</g>\n",
       "<!-- _zone_counts_inputs -->\n",
       "<g id=\"node20\" class=\"node\">\n",
       "<title>_zone_counts_inputs</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"301.45,-381.3 144.35,-381.3 144.35,-336.7 301.45,-336.7 301.45,-381.3\"/>\n",
       "<text text-anchor=\"start\" x=\"159.15\" y=\"-353.2\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">aggregation_level</text>\n",
       "<text text-anchor=\"start\" x=\"271.65\" y=\"-353.2\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">str</text>\n",
       "</g>\n",
       "<!-- _zone_counts_inputs&#45;&gt;zone_counts -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>_zone_counts_inputs&#45;&gt;zone_counts</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M301.74,-371.24C311.39,-372.38 321.13,-373.35 330.45,-374 416.35,-380.02 819.73,-375.5 887.75,-370 895.49,-369.37 903.6,-368.47 911.6,-367.43\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"911.99,-370.78 921.42,-365.94 911.03,-363.84 911.99,-370.78\"/>\n",
       "</g>\n",
       "<!-- _zone_counts_inputs&#45;&gt;level_info -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>_zone_counts_inputs&#45;&gt;level_info</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M301.79,-346.93C324.19,-343.45 348.15,-339.74 368.86,-336.53\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"369.34,-339.84 378.69,-334.85 368.27,-332.92 369.34,-339.84\"/>\n",
       "</g>\n",
       "<!-- input -->\n",
       "<g id=\"node21\" class=\"node\">\n",
       "<title>input</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"84.68,-464.3 30.68,-464.3 30.68,-427.7 84.68,-427.7 84.68,-464.3\"/>\n",
       "<text text-anchor=\"middle\" x=\"57.68\" y=\"-440.2\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">input</text>\n",
       "</g>\n",
       "<!-- function -->\n",
       "<g id=\"node22\" class=\"node\">\n",
       "<title>function</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M80.1,-409.3C80.1,-409.3 35.25,-409.3 35.25,-409.3 29.25,-409.3 23.25,-403.3 23.25,-397.3 23.25,-397.3 23.25,-384.7 23.25,-384.7 23.25,-378.7 29.25,-372.7 35.25,-372.7 35.25,-372.7 80.1,-372.7 80.1,-372.7 86.1,-372.7 92.1,-378.7 92.1,-384.7 92.1,-384.7 92.1,-397.3 92.1,-397.3 92.1,-403.3 86.1,-409.3 80.1,-409.3\"/>\n",
       "<text text-anchor=\"middle\" x=\"57.68\" y=\"-385.2\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">function</text>\n",
       "</g>\n",
       "<!-- cluster -->\n",
       "<g id=\"node23\" class=\"node\">\n",
       "<title>cluster</title>\n",
       "<path fill=\"#ffffff\" stroke=\"#649fb3\" stroke-dasharray=\"5,2\" d=\"M76.35,-354.3C76.35,-354.3 39,-354.3 39,-354.3 33,-354.3 27,-348.3 27,-342.3 27,-342.3 27,-329.7 27,-329.7 27,-323.7 33,-317.7 39,-317.7 39,-317.7 76.35,-317.7 76.35,-317.7 82.35,-317.7 88.35,-323.7 88.35,-329.7 88.35,-329.7 88.35,-342.3 88.35,-342.3 88.35,-348.3 82.35,-354.3 76.35,-354.3\"/>\n",
       "<text text-anchor=\"middle\" x=\"57.68\" y=\"-330.2\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">cluster</text>\n",
       "</g>\n",
       "<!-- field -->\n",
       "<g id=\"node24\" class=\"node\">\n",
       "<title>field</title>\n",
       "<path fill=\"#c8dae0\" stroke=\"black\" d=\"M77.1,-300.05C77.1,-300.05 38.25,-300.05 38.25,-300.05 32.25,-300.05 26.25,-294.05 26.25,-288.05 26.25,-288.05 26.25,-273.95 26.25,-273.95 26.25,-267.95 32.25,-261.95 38.25,-261.95 38.25,-261.95 77.1,-261.95 77.1,-261.95 83.1,-261.95 89.1,-267.95 89.1,-273.95 89.1,-273.95 89.1,-288.05 89.1,-288.05 89.1,-294.05 83.1,-300.05 77.1,-300.05\"/>\n",
       "<text text-anchor=\"middle\" x=\"57.68\" y=\"-275.95\" font-family=\"Courier,monospace\" font-size=\"14.00\">field</text>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x29d00fe10>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "aggregation_level = \"avatarId\"\n",
    "zone_features = importlib.import_module(\"features\")\n",
    "\n",
    "dr = driver.Driver({}, zone_features, adapter=adapter)  # can pass in multiple modules\n",
    "dr.display_all_functions(deduplicate_inputs=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute the features of interest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "feature_groups_of_interest = [\"level_info\", \"zone_counts\"]\n",
    "features = dr.execute(feature_groups_of_interest, inputs={\"aggregation_level\": aggregation_level})\n",
    "features.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
